# -*- coding:utf-8 -*-
"""
# @file name    : transforms_method_1.py
# @author       : QuZhang
# @date         : 2020-12-05 23:10
# @brief        : transforms方法(一)
"""
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from torchvision import transforms
from tools.my_dataset import RMBDataset
from torch.utils.data import DataLoader
import torch
from PIL import Image
import matplotlib.pyplot as plt
from tools.common_tools import transform_invert


BATCH_SIZE = 4

if __name__ == "__main__":

    split_dir = os.path.abspath(os.path.join(BASE_DIR, "..", '..', 'data', 'rmb_split'))
    if not os.path.exists(split_dir):
        raise Exception(r"数据{} 不存在，回到lesson-06-1_split_dataset.py生成数据".format(split_dir))
    train_dir = os.path.join(split_dir, 'train')
    norm_mean = [0.485, 0.456, 0.406]
    norm_std = [0.229, 0.224, 0.225]
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),

        # 1. 对图像边缘进行填充
        # transforms.Pad(padding=32, fill=(255, 0, 0), padding_mode='constant'),
        # transforms.Pad(padding=(8, 64), fill=(255, 0, 0), padding_mode="constant"),
        # transforms.Pad(padding=(8, 10, 20, 32), fill=(255, 0, 0)),
        # transforms.Pad(padding=(8, 16, 32, 64), fill=(255, 0, 0), padding_mode='symmetric'),

        # 2. 实用方法：调整亮度、对比度、饱和度和色相
        # transforms.ColorJitter(brightness=0.5),
        # transforms.ColorJitter(contrast=0.5),
        # transforms.ColorJitter(saturation=0.5),
        # transforms.ColorJitter(hue=0.3),

        # 3. 变为灰度图
        # transforms.Grayscale(num_output_channels=3),  # 一定会执行灰度
        # transforms.RandomGrayscale(p=0.1),

        # 4. 对图像进行仿射变换:旋转、平移、缩放、错切、翻转
        # transforms.RandomAffine(degrees=30),  # 旋转30度
        # transforms.RandomAffine(degrees=0, translate=(0.2, 0.2), fillcolor=(255, 0, 0)),  # 平移，之后填充颜色
        # transforms.RandomAffine(degrees=0, scale=(0.7, 0.7)),  # 缩小30%,只为原来的70%
        # transforms.RandomAffine(degrees=0, shear=(0, 0, 0, 45)),  # 错切

        # 实用数据增强
        # 5 随机遮挡
        # transforms.ToTensor(),  # 随机遮挡是在张量上进行操作
        # transforms.RandomErasing(p=1, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=(254/255, 0, 0)),  # 指定遮挡的像素值
        # transforms.RandomErasing(p=1, scale=(0.02, 0.33), ratio=(0.3, 3.3), value="random"),  # 遮挡区域的像素值使用随机采样

        # 选择一组方法进行数据增强
        # 1 选择一个进行执行
        # transforms.RandomChoice([transforms.RandomVerticalFlip(p=1), transforms.RandomHorizontalFlip(p=1)]),

        # 2 根据概率执行一组
        # transforms.RandomApply([transforms.RandomAffine(degrees=0, shear=45, fillcolor=(255, 0, 0)),
        #                         transforms.Grayscale(num_output_channels=3)], p=0.5),

        # 3 随机对一组变换排序后再执行
        transforms.RandomOrder([transforms.RandomRotation(15),
                                transforms.Pad(padding=32, fill=(255, 0, 0)),
                                transforms.RandomAffine(degrees=0, translate=(0.01, 0.1), scale=(0.9, 1.1), fillcolor=(0, 255, 0))]),

        transforms.ToTensor(),
        transforms.Normalize(norm_mean, norm_std),
    ])

    train_dir = RMBDataset(data_dir=train_dir, transform=train_transform)

    train_loader = DataLoader(dataset=train_dir, batch_size=BATCH_SIZE, shuffle=True)

    for i, data in enumerate(train_loader):

        if i < 1:
            # 批量加载的图片一般是4维的张量(由多个三维张量构成)：图片张数，每张图片是一个三维张量,每张图片的每个通道是一个二维张量
            inputs, labels = data  # inputs:(B, C, H, W)
            # print("img_tensor:{}".format(inputs))
            print("label: {}".format(labels))

            for j in range(BATCH_SIZE):
                # 获取每一种图片对应的张量,根据四维张量的第一维度一个一个读取图片对应的三维张量
                img_tensor = inputs[j, ...]  # (C, H, W)
                # print("img_tensor:", img_tensor)
                img = transform_invert(img_tensor, train_transform)
                plt.imshow(img)
                plt.show()
                plt.pause(0.5)
                plt.close()
